#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/20
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
import torch
import torch.optim as optim
import itertools
from six_three_build_network import ReplayBuffer
from six_three_build_network import Generator
from six_three_build_network import Discriminator
from six_two_load_data import _get_train_data
from six_two_load_data import to_img
from six_two_load_data import save_image

# 代码6-6
# 训练网络
# 创建两个用于存储生成图像的缓冲区
fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()

# 创建两个生成器，一个用于A到B的转换，一个用于B到A的转换
netG_A2B = Generator()
netG_B2A = Generator()
# 创建两个判别器，一个用于A的判别，一个用于B的判别
netD_A = Discriminator()
netD_B = Discriminator()

# 设置生成器的损失函数为均方误差损失
criterion_GAN = torch.nn.MSELoss()
# 设置循环损失函数为L1损失
criterion_cycle = torch.nn.L1Loss()
# 设置恒等损失函数为L1损失
criterion_identity = torch.nn.L1Loss()

# 设置判别器的学习率为3e-4
d_learning_rate = 3e-4
# 设置生成器的学习率为3e-4
g_learning_rate = 3e-4
# 设置优化器的beta参数为(0.5, 0.999)
optim_betas = (0.5, 0.999)

# 定义生成器G的优化器，使用Adam优化算法，优化参数包括netG_A2B和netG_B2A的所有参数
g_optimizer = optim.Adam(itertools.chain(netG_A2B.parameters(), netG_B2A.parameters()),
                         lr=d_learning_rate)
# 定义判别器D_A的优化器，使用Adam优化算法，优化参数为netD_A的所有参数
da_optimizer = optim.Adam(netD_A.parameters(), lr=d_learning_rate)
# 定义判别器D_B的优化器，使用Adam优化算法，优化参数为netD_B的所有参数
db_optimizer = optim.Adam(netD_B.parameters(), lr=d_learning_rate)

num_epochs = 1000
# 设定批处理大小为1
batch_size = 1

for epoch in range(num_epochs):

    # 获取训练数据，返回真实数据real_a和真实数据real_b
    real_a, real_b = _get_train_data(batch_size)
    # 创建一个全为1的tensor，长度为batch_size，用于生成真实数据的目标值
    target_real = torch.full((batch_size,), 1).float()
    # 创建一个全为0的tensor，长度为batch_size，用于生成假数据的目标值
    target_fake = torch.full((batch_size,), 0).float()

    # 清空梯度，准备计算新的梯度
    g_optimizer.zero_grad()

    # 第一步：训练生成器
    # 使用netG_A2B网络将真实的B域图像转化为A域，并将其转化为浮点数
    same_B = netG_A2B(real_b).float()
    # 计算转化后的图像与原始图像之间的差异，并乘以5.0作为损失
    loss_identity_B = criterion_identity(same_B, real_b) * 5.0
    # 使用netG_B2A网络将真实的A域图像转化为B域，并将其转化为浮点数
    same_A = netG_B2A(real_a).float()
    # 计算转化后的图像与原始图像之间的差异，并乘以5.0作为损失
    loss_identity_A = criterion_identity(same_A, real_a) * 5.0

    # 使用netG_A2B生成器将real_a转换为fake_B
    fake_B = netG_A2B(real_a).float()
    # 使用netD_B判别器对fake_B进行判别
    pred_fake = netD_B(fake_B).float()
    # 计算A到B的生成器GAN损失
    loss_GAN_A2B = criterion_GAN(pred_fake, target_real)
    # 使用netG_B2A生成器将real_b转换为fake_A
    fake_A = netG_B2A(real_b).float()
    # 使用netD_A判别器对fake_A进行判别
    pred_fake = netD_A(fake_A).float()
    # 计算B到A的生成器GAN损失
    loss_GAN_B2A = criterion_GAN(pred_fake, target_real)
    # 使用netG_B2A生成器将fake_B恢复为recovered_A
    recovered_A = netG_B2A(fake_B).float()
    # 计算A到B的循环损失
    loss_cycle_ABA = criterion_cycle(recovered_A, real_a) * 10.0
    # 使用netG_A2B生成器将fake_A恢复为recovered_B
    recovered_B = netG_A2B(fake_A).float()
    # 计算B到A的循环损失
    loss_cycle_BAB = criterion_cycle(recovered_B, real_b) * 10.0
    # 计算总的生成器损失
    loss_G = (loss_identity_A + loss_identity_B + loss_GAN_A2B +
              loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB)
    # 反向传播计算梯度
    loss_G.backward()
    # 更新生成器的参数
    g_optimizer.step()

    # 第二步：训练判别器
    # 训练判别器A
    da_optimizer.zero_grad()  # 清空优化器梯度
    pred_real = netD_A(real_a).float()  # 使用网络netD_A对真实图像进行预测，并转换为浮点数
    loss_D_real = criterion_GAN(pred_real, target_real)  # 计算真实图像的损失
    fake_A = fake_A_buffer.push_and_pop(fake_A)  # 将生成的假图像添加到缓存中，并可能从缓存中弹出一些旧的假图像
    pred_fake = netD_A(fake_A.detach()).float()  # 对假图像进行预测，但不更新生成的假图像，因为它已经在缓存中
    loss_D_fake = criterion_GAN(pred_fake, target_fake)  # 计算假图像的损失
    loss_D_A = (loss_D_real + loss_D_fake) * 0.5  # 计算总的损失
    loss_D_A.backward()  # 反向传播计算梯度
    da_optimizer.step()  # 更新优化器参数
    # 训练判别器B
    db_optimizer.zero_grad()  # 清空梯度，准备计算新的梯度
    pred_real = netD_B(real_b)  # 使用网络netD_B对真实图片进行预测
    loss_D_real = criterion_GAN(pred_real, target_real)  # 计算真实图片的损失
    fake_B = fake_B_buffer.push_and_pop(fake_B)  # 将生成的假图片加入缓存并弹出旧的假图片
    pred_fake = netD_B(fake_B.detach())  # 对假图片进行预测，但不更新生成器参数
    loss_D_fake = criterion_GAN(pred_fake, target_fake)  # 计算假图片的损失
    loss_D_B = (loss_D_real + loss_D_fake) * 0.5  # 计算总的损失
    loss_D_B.backward()  # 反向传播计算梯度
    db_optimizer.step()  # 更新优化器的参数

    # 损失打印，存储伪造图片
    print('Epoch[{}],loss_G:{:.6f} ,loss_D_A:{:.6f},loss_D_B:{:.6f}'
          .format(epoch, loss_G.data.item(), loss_D_A.data.item(),
                  loss_D_B.data.item()))
    if (epoch + 1) % 20 == 0 or epoch == 0:
        b_fake = to_img(fake_B.data)
        a_fake = to_img(fake_A.data)
        a_real = to_img(real_a.data)
        b_real = to_img(real_b.data)
        save_image(a_fake, 'tmp/6-a_fake.png')
        save_image(b_fake, 'tmp/6-b_fake.png')
        save_image(a_real, 'tmp/6-a_real.png')
        save_image(b_real, 'tmp/6-b_real.png')


